#!/usr/bin/env python3

"""
主窗口底部是“打开摄像头”按钮，按钮上面显示照片路径列表。

单击“打开摄像头”，弹出摄像头窗口，开一个单独线程在窗口左边显示摄像头实时画面，下边是拍照按钮，右边显示拍下的照片。

单击“拍照”会把摄像头画面的当前帧保存到本地、显示到右侧照片列表。单击照片下面的删除按钮会从本地删掉它。新增、删除照片会实时更新主窗口的路径列表。照片拍满 5 张会禁用“拍照”按钮，删除照片会取消禁用。

单击照片弹出放大窗口，在放大窗口单击按钮或 CTRL+鼠标滚轮会缩放照片。

摄像头窗口在关闭时不会清空右侧照片，重新打开时仍显示之前拍的照片。关闭时会暂停获取摄像头画面。
"""

import os
import os.path as os_path
import random
import string
import sys
from datetime import datetime
from typing import Callable, Generator, Optional, Tuple

import cv2
from cv2.typing import MatLike
from PySide6.QtCore import QEvent, QModelIndex, QThread, Signal
from PySide6.QtGui import (
    QCursor,
    QImage,
    QMouseEvent,
    QPixmap,
    QStandardItem,
    QStandardItemModel,
    Qt,
    QWheelEvent,
)
from PySide6.QtWidgets import (
    QApplication,
    QDialog,
    QGridLayout,
    QLabel,
    QListView,
    QMessageBox,
    QPushButton,
    QScrollArea,
    QScrollBar,
    QSizePolicy,
    QTableView,
    QVBoxLayout,
    QWidget,
)


class Dialog_Photo(QDialog):
    def __init__(self, pixmap: QPixmap, parent=None):
        super().__init__(parent)

        self.label_photo = QLabel()
        self.label_photo.setSizePolicy(QSizePolicy.Policy.Ignored, QSizePolicy.Policy.Ignored)
        # When enabled and the label shows a pixmap, it will scale the pixmap to fill the available space.
        self.label_photo.setScaledContents(True)
        self.label_photo.setPixmap(pixmap)

        self.scroll_area = QScrollArea()
        self.scroll_area.setWidget(self.label_photo)

        self.btn_zoom_in = QPushButton("放大")
        self.btn_zoom_in.clicked.connect(self.zoom_in)
        self.btn_zoom_out = QPushButton("缩小")
        self.btn_zoom_out.clicked.connect(self.zoom_out)
        self.btn_normal_size = QPushButton("原始尺寸")
        self.btn_normal_size.clicked.connect(self.normal_size)

        layout_grid = QGridLayout()
        colno = 0
        for colno, btn in enumerate((self.btn_zoom_in, self.btn_zoom_out, self.btn_normal_size)):
            layout_grid.addWidget(btn, 1, colno)
        layout_grid.addWidget(self.scroll_area, 0, 0, 1, colno + 1)
        self.setLayout(layout_grid)

        self.scale_factor = 1.0
        self.max_scale_factor = 4.0
        self.min_scale_factor = 0.25

    def wheelEvent(self, event: QWheelEvent) -> None:
        """
        CTRL+滚轮缩放图片
        """
        if event.modifiers() == Qt.KeyboardModifier.ControlModifier:
            if event.angleDelta().y() > 0 and self.scale_factor < self.max_scale_factor:
                self.zoom_in()
            elif event.angleDelta().y() < 0 and self.scale_factor > self.min_scale_factor:
                self.zoom_out()
        return super().wheelEvent(event)

    def zoom_in(self):
        self.scale_photo(1.25)

    def zoom_out(self):
        self.scale_photo(0.8)

    def normal_size(self):
        self.label_photo.adjustSize()
        self.scale_factor = 1.0

    def scale_photo(self, factor: float):
        self.scale_factor *= factor
        self.label_photo.resize(self.label_photo.pixmap().size() * self.scale_factor)

        self.adjust_scrollbar(self.scroll_area.horizontalScrollBar(), factor)
        self.adjust_scrollbar(self.scroll_area.verticalScrollBar(), factor)

        self.btn_zoom_in.setEnabled(self.scale_factor < self.max_scale_factor)
        self.btn_zoom_out.setEnabled(self.scale_factor > self.min_scale_factor)

    def adjust_scrollbar(self, scroll_bar: QScrollBar, factor: float):
        scroll_bar.setValue(int(factor * scroll_bar.value() + ((factor - 1) * scroll_bar.pageStep() / 2)))


class Label_Camera(QLabel):
    def __init__(self, parent=None):
        super().__init__(parent)

    def update_frame(self, frame: MatLike):
        """
        被 Thread_Camera 调用来持续更新画面
        """
        pixmap = self._frame2pixmap(frame)
        self.setPixmap(pixmap)

    def _frame2pixmap(self, frame: MatLike) -> QPixmap:
        """
        把 OpenCV 获取的 frame 转换成 QPixmap 并返回。QPixmap 显示到 QLabel 上，预览照片。
        """
        h, w, ch = frame.shape
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        bytes_per_line = ch * w

        img = QImage(frame_rgb.data, w, h, bytes_per_line, QImage.Format.Format_RGB888)
        return QPixmap.fromImage(img)


class Label_Photo(QLabel):
    """
    摄像头对话框右侧的照片，用来预览。单击时弹窗显示放大的照片。

    photo_path (str): 照片的保存路径，但这个class并不负责保存操作。这个参数只在用户单击照片底部的"删除"按钮时被TableView_Photo获取然后从本地删除该路径下的文件。保存操作会在这个class实例化之前被Dialog_Camera执行。
    pixmap (QPixmap): 从左侧摄像头画面获取的QPixmap。
    """

    def __init__(
        self,
        photo_path: str,
        pixmap: QPixmap,
        *,
        scale: Optional[Tuple[int, int]] = None,
        parent=None,
    ):
        super().__init__(parent)
        self.photo_path = photo_path

        self.set_scaled_pixmap(pixmap, scale=scale)
        self.pixmap_orig = pixmap

    def set_scaled_pixmap(self, pixmap: QPixmap, *, scale: Optional[Tuple[int, int]] = None):
        if scale is not None:
            pixmap = pixmap.scaled(*scale, Qt.AspectRatioMode.KeepAspectRatio)
        self.setPixmap(pixmap)

    def mousePressEvent(self, event: QMouseEvent):
        """
        单击时弹窗显示放大的照片
        """
        if event.button() == Qt.MouseButton.LeftButton and not self.pixmap_orig.isNull():
            Dialog_Photo(self.pixmap_orig, self).open()
        super().mousePressEvent(event)


class Model_Photo(QStandardItemModel):
    """
    rows (int): 最大允许的拍照数量
    """
    def __init__(self, rows: int, parent=None):
        super().__init__(rows, 1, parent)

    def empty(self):
        rowno = self.rowCount()
        colno = self.columnCount()

        self.clear()
        self.setRowCount(rowno)
        self.setColumnCount(colno)


class TableView_Photo(QTableView):
    # 拍照达到最大允许的照片数量时发信号，Dialog_Camera 接收后，禁用拍照按钮
    can_take_more_photos = Signal(bool)
    # 删除照片后发信号，传递行号，Widget_Camera 接收后从路径列表删除对应行
    deletion_request = Signal(int)

    def __init__(self, parent=None):
        super().__init__(parent)

        self.horizontalHeader().hide()
        self.verticalHeader().hide()
        self.setShowGrid(False)
        self.setFocusPolicy(Qt.FocusPolicy.NoFocus)
        self.setEditTriggers(QTableView.EditTrigger.NoEditTriggers)
        self.setSelectionMode(QTableView.SelectionMode.NoSelection)
        self.setHorizontalScrollMode(QTableView.ScrollMode.ScrollPerPixel)
        self.setVerticalScrollMode(QTableView.ScrollMode.ScrollPerPixel)

        self.next_row: int = 0

    def _attach_delete_btn(self, label_photo: QLabel) -> QVBoxLayout:
        """
        接收一个用来预览照片的 QLabel，添加删除按钮
        返回 QVBoxLayout
        """
        btn_delete = QPushButton("删除")
        btn_delete.clicked.connect(self.delete_photo)
        layout_photo = QVBoxLayout()
        layout_photo.addWidget(label_photo)
        layout_photo.addWidget(btn_delete)

        return layout_photo

    def append_photo(self, label_photo: Label_Photo):
        layout_photo = self._attach_delete_btn(label_photo)
        widget_photo = QWidget()
        widget_photo.setLayout(layout_photo)

        idx = self.model().index(self.next_row, 0)
        self.setIndexWidget(idx, widget_photo)

        self.resizeRowsToContents()
        self.resizeColumnsToContents()
        self.scrollToBottom()

        self.next_row += 1
        # self.model().rowCount() 是在 model 初始化时设置的
        if self.next_row + 1 > self.model().rowCount():
            self.can_take_more_photos.emit(False)

    def is_cell_empty(self, idx: QModelIndex) -> bool:
        widget = self.indexWidget(idx)
        return widget is None

    def yield_photo_path(self) -> Generator[str, None, None]:
        for rowno in range(self.model().rowCount()):
            for colno in range(self.model().columnCount()):
                idx = self.model().index(rowno, colno)
                if not self.is_cell_empty(idx):
                    yield self.get_photo_path(idx)

    def get_photo_path(self, idx: QModelIndex) -> str:
        widget = self.indexWidget(idx)
        label_photo = widget.layout().itemAt(0).widget()
        if not isinstance(label_photo, Label_Photo):
            raise TypeError
        photo_path = label_photo.photo_path
        return photo_path

    def delete_photo(self):
        cursor_pos = self.mapFromGlobal(QCursor.pos())
        idx = self.indexAt(cursor_pos)
        try:
            photo_path = self.get_photo_path(idx)
        except TypeError:
            QMessageBox.critical(self, "Error", "无法获取保存路径，未删除")
            return

        # 删除本地照片
        if os_path.exists(photo_path):
            os.remove(photo_path)

        # 删除预览界面照片
        self.model().removeRow(idx.row())
        self.model().setRowCount(self.model().rowCount() + 1)

        # 发信号给 Widget_Camera
        self.deletion_request.emit(idx.row())

        self.next_row -= 1
        self.can_take_more_photos.emit(True)


class Dialog_Camera(QDialog):
    # 拍照后发信号，传递照片路径，Widget_Camera 接收后添加到路径列表
    addition_request = Signal(str)

    def __init__(self, parent=None):
        super().__init__(parent)

        self.init_ui()
        self.init_camera()

    def init_ui(self):
        self.resize(960, 540)

        layout_dialog = QGridLayout()
        self.setLayout(layout_dialog)

        self.label_camera = Label_Camera()
        self.btn_capture = QPushButton("拍照")
        self.btn_capture.clicked.connect(self.on_capture)

        self.model_preview = Model_Photo(5, self)  # 最多拍 5 张
        self.tableview_photo = TableView_Photo(self)
        self.tableview_photo.setModel(self.model_preview)
        self.tableview_photo.can_take_more_photos.connect(self.btn_capture.setEnabled)

        layout_dialog.addWidget(self.label_camera, 0, 0)
        layout_dialog.addWidget(self.btn_capture, 1, 0, 1, 2)
        layout_dialog.addWidget(self.tableview_photo, 0, 1)

        self.finished.connect(self.on_finished)

    def init_camera(self):
        self.cap = cv2.VideoCapture(0)
        if not self.cap.isOpened():
            QMessageBox.critical(self, "Error", "无法连接摄像头")
            exit()

        self.thread_camera = Thread_Camera(self.cap, self.label_camera.update_frame)

    def _generate_photo_path(self) -> str:
        """
        生成保存路径
        例如 pictures/2024/202404/20240402/20240402xxxxxx.jpg, xxxxxx 为 6 位随机大写字母
        """
        now = datetime.now()
        year = now.strftime("%Y")
        month = now.strftime("%Y%m")
        day = now.strftime("%Y%m%d")

        directory = f"pictures/{year}/{month}/{day}"
        if not os.path.exists(directory):
            os.makedirs(directory)

        random_photo_path = "".join(random.choices(string.ascii_uppercase, k=6))
        photo_path = f"{directory}/{day}{random_photo_path}.jpg"
        return photo_path

    def on_capture(self):
        """
        1. 从左侧摄像头画面获取 pixmap
        2. 把照片保存到本地
        3. 把照片发送到右侧预览区
        4. 发信号给 Widget_Camera 让更新照片路径列表
        """
        photo_path = self._generate_photo_path()
        pixmap = self.label_camera.pixmap()
        pixmap.save(photo_path)

        label_photo = Label_Photo(photo_path, pixmap=pixmap, scale=(100, 100))
        self.tableview_photo.append_photo(label_photo)

        self.addition_request.emit(photo_path)

    def open(self):
        self.thread_camera.start()
        return super().open()

    def on_finished(self) -> None:
        self.thread_camera.stop()


class Thread_Camera(QThread):
    def __init__(self, cap: cv2.VideoCapture, update_frame: Callable):
        QThread.__init__(self)
        self.cap = cap
        self.update_frame = update_frame
        self.running = True

    def run(self):
        while self.running:
            ok, frame = self.cap.read()
            if not ok or frame is None:
                continue

            self.update_frame(frame)

    def start(self):
        self.running = True
        return super().start()

    def stop(self):
        self.running = False


class Widget_Camera(QWidget):
    def __init__(self):
        super().__init__()

        self.init_ui()

    def init_ui(self):
        layout_grid = QVBoxLayout()
        self.setLayout(layout_grid)

        self.dialog_camera = Dialog_Camera(self)
        self.dialog_camera.tableview_photo.deletion_request.connect(self.on_photo_deletion)
        self.dialog_camera.addition_request.connect(self.on_photo_addition)

        self.model_photo_paths = QStandardItemModel(self)
        self.model_photo_paths.setHorizontalHeaderLabels(("照片路径",))
        self.listview_photo_paths = QListView(self)
        self.listview_photo_paths.setModel(self.model_photo_paths)

        btn_launch_camera = QPushButton("打开摄像头")
        btn_launch_camera.clicked.connect(self.dialog_camera.open)

        layout_grid.addWidget(self.listview_photo_paths)
        layout_grid.addWidget(btn_launch_camera)

    def on_photo_deletion(self, rowno: int):
        self.model_photo_paths.removeRow(rowno)

    def on_photo_addition(self, photo_path: str):
        self.model_photo_paths.appendRow(QStandardItem(photo_path))


if __name__ == "__main__":
    app = QApplication([])
    widget = Widget_Camera()
    widget.show()
    sys.exit(app.exec())
